import gymnasium as gym
from gymnasium.spaces import Dict, Graph, Tuple, Discrete, Box, MultiDiscrete, GraphInstance
import numpy as np

from src.utils import OperationsGraph

TIME_MAX = 1e3

class env_FJSP(gym.Env):

    metadata = {"render_modes": ["console"]}

    def __init__(self, problem:OperationsGraph) -> None:
        self.problem = problem
        self.observation_space = Dict({
            # "operations_graph": Graph(
            #     node_space = Box(low=-1.0, high=1e3, shape=problem.graph_np[0].shape)  # (作业数, 工序数, 机器数, 加工时间)
            # )
            "waiting_operations": Box(low=-1.0, high=problem.operation_count, shape=([problem.jobs, problem.machines, 4] )) # 等待上机的操作
            ,"machines_status": Dict({
                "current_opid" : Tuple([Discrete(problem.operation_count, start=-1)] * problem.machines),  # 几个机器当前加工的是哪个工序
                # "healthiness": Box(0.0, 1e3, shape=(problem.machines, )),   # 机器健康度
                "finish_time": Box(0.0, TIME_MAX, shape=(problem.machines, )),   # 机器加工的完成时间
            })
            ,"current_time": Box(0, TIME_MAX)
        })

        self.action_space = Tuple([Discrete(problem.machines, start=-1)] * problem.machines)


    def step(self, action):
        self._current_time += 1.0 
        idle_machines = self._get_idle_machines()
        for i in range(self.problem.machines):
            if idle_machines[i] :
                last_opid = self._current_opids[i]
                self._current_opids[i] = action[i]
                self._finish_times[i] = self._waiting_operations[action[i]][i][3] + self._current_time
                # self._graph.edge_links.append([last_opid, action[i]])

        observation = self._get_obs()
        reward = self._calculate_reward()
        terminated = self._current_time > TIME_MAX
        truncated = False
        info = self._get_info()

        return observation, reward, terminated, truncated, info
    
    def reset(self, seed=None, options=None):
        # self.problem.reset()
        self._graph = GraphInstance(self.problem.graph_np, edges=None, edge_links=self.problem.edges)
        self._current_opids = np.zeros(self.problem.machines, dtype=np.int32)
        self._finish_times = np.zeros(self.problem.machines, dtype=np.float32)
        # self._current_time = np.array([0.0, ], dtype=np.float32)
        self._current_time = 0.0
        self._waiting_operations = np.array([self._graph.nodes[_[1]] for _ in self._graph.edge_links if _[0] == 0], dtype=np.float32)

        observation = self._get_obs()
        info = self._get_info()
        print(observation["machines_status"])
        return observation, info
    
    def render(self):
        ...

    def close(self):
        ...

    def _calculate_reward(self):
        # 根据问题需求计算奖励
        return 1.0
    
    def _get_info(self):
        # 返回环境的额外信息
        return {}


    def _get_obs(self):
        return {
            # "operations_graph": self._graph
            "waiting_operations": self._waiting_operations
            ,"machines_status": {
                "current_opid": self._current_opids,  # 几个机器当前加工的是哪个工序
                "finish_time": self._finish_times,   # 机器加工的完成时间
            }
            ,"current_time": np.array([self._current_time, ], dtype=np.float32)
        }
    
    def _get_idle_machines(self):
        """有哪些机器空闲"""
        return self._finish_times <= self._current_time
    
if __name__ == "__main__":
    from stable_baselines3.common.env_checker import check_env
    env = env_FJSP(OperationsGraph("./Mk01.fjs"))
    # If the environment don't follow the interface, an error will be thrown
    check_env(env, warn=True)